import os

import pandas as pd


class DocumentLoader:
    def __init__(self, data_dir="../dataset"):
        self.data_dir = data_dir
        self.data_files = [
            "(sample)sam_tianchi_mum_baby.csv",
            "(sample)sam_tianchi_mum_baby_trade_history.csv",
            "customer_shopping_data.csv",
            "products.csv",
            "优衣库销售数据.xlsx",
            "商城详细销售数据.xlsx"
        ]

    def fetch_ecommerce_docs(self):
        """从本地数据集提取电商数据文档"""
        docs = []

        # 处理母婴数据集
        if "(sample)sam_tianchi_mum_baby.csv" in self.data_files:
            doc = self._process_mother_baby_data()
            if doc: docs.append(doc)

        # 处理购物交易数据
        if "(sample)sam_tianchi_mum_baby_trade_history.csv" in self.data_files:
            doc = self._process_trade_history()
            if doc: docs.append(doc)

        # 处理客户购物数据
        if "customer_shopping_data.csv" in self.data_files:
            doc = self._process_customer_shopping()
            if doc: docs.append(doc)

        # 处理产品目录数据
        if "products.csv" in self.data_files:
            doc = self._process_product_catalog()
            if doc: docs.append(doc)

        # 处理优衣库销售数据
        if "优衣库销售数据.xlsx" in self.data_files:
            doc = self._process_uniqlo_sales()
            if doc: docs.append(doc)

        # 处理商城销售数据
        if "商城详细销售数据.xlsx" in self.data_files:
            doc = self._process_mall_sales()
            if doc: docs.append(doc)

        return docs

    def _load_dataframe(self, filename):
        """加载指定数据文件为DataFrame"""
        filepath = os.path.join(self.data_dir, filename)

        if filename.endswith('.csv'):
            return pd.read_csv(filepath)
        elif filename.endswith('.xlsx'):
            return pd.read_excel(filepath)
        return None

    def _create_doc(self, source, title, data_desc, sample_data):
        """创建标准格式文档"""
        return {
            "source": f"file://{source}",
            "content": f"""
                {title}
                ====================
                数据描述:
                {data_desc}

                数据样本:
                {sample_data}
                """
        }

    def _process_mother_baby_data(self):
        """处理母婴用户数据"""
        df = self._load_dataframe("(sample)sam_tianchi_mum_baby.csv")
        if df is None or df.empty:
            return None

        data_desc = """
            - 包含950位母婴用户的基本信息
            - 数据字段: 用户ID, 出生日期, 性别, 城市等
            - 涵盖怀孕阶段和婴幼儿阶段用户数据
            """
        sample = df.head(3).to_string(index=False)

        return self._create_doc(
            source="(sample)sam_tianchi_mum_baby.csv",
            title="母婴用户画像数据",
            data_desc=data_desc,
            sample_data=sample
        )

    def _process_trade_history(self):
        """处理母婴交易历史数据"""
        df = self._load_dataframe("(sample)sam_tianchi_mum_baby_trade_history.csv")
        if df is None or df.empty:
            return None

        data_desc = """
            - 包含2.9万条母婴产品交易记录
            - 数据字段: 用户ID, 商品ID, 类别ID, 购买日期等
            - 时间跨度: 3年的交易数据
            """
        sample = df.head(3).to_string(index=False)

        return self._create_doc(
            source="(sample)sam_tianchi_mum_baby_trade_history.csv",
            title="母婴商品交易记录",
            data_desc=data_desc,
            sample_data=sample
        )

    def _process_customer_shopping(self):
        """处理客户购物数据"""
        df = self._load_dataframe("customer_shopping_data.csv")
        if df is None or df.empty:
            return None

        data_desc = """
            - 包含客户购物行为数据
            - 数据字段: 客户ID, 购物日期, 品类, 金额等
            - 涉及多个产品类别的购物记录
            """
        sample = df.head(3).to_string(index=False)

        return self._create_doc(
            source="customer_shopping_data.csv",
            title="客户购物行为数据集",
            data_desc=data_desc,
            sample_data=sample
        )

    def _process_product_catalog(self):
        """处理产品目录数据"""
        df = self._load_dataframe("products.csv")
        if df is None or df.empty:
            return None

        data_desc = """
            - 商品基础信息数据
            - 数据字段: 产品ID, 产品名称, 类别, 价格等
            - 多品类商品数据
            """
        sample = df.head(3).to_string(index=False)

        return self._create_doc(
            source="products.csv",
            title="商品目录数据集",
            data_desc=data_desc,
            sample_data=sample
        )

    def _process_uniqlo_sales(self):
        """处理优衣库销售数据"""
        df = self._load_dataframe("优衣库销售数据.xlsx")
        if df is None or df.empty:
            return None

        # 分析数据格式和内容
        data_desc = """
            - 优衣库产品销售数据
            - 包含: 日期, 产品系列, 尺码, 颜色, 销售量等
            - 分类统计各产品线销售表现
            """

        # 只显示关键列的前几行
        key_columns = df.columns[:5] if len(df.columns) > 5 else df.columns
        sample = df[key_columns].head(3).to_string(index=False)

        return self._create_doc(
            source="优衣库销售数据.xlsx",
            title="优衣库产品销售数据",
            data_desc=data_desc,
            sample_data=sample
        )

    def _process_mall_sales(self):
        """处理商城销售数据"""
        df = self._load_dataframe("商城详细销售数据.xlsx")
        if df is None or df.empty:
            return None

        data_desc = """
            - 商场详细交易记录
            - 包括: 交易ID, 销售日期, 店铺ID, 产品ID等
            - 含销售时间、金额、折扣等信息
            """
        sample = df.head(3).to_string(index=False)

        return self._create_doc(
            source="商城详细销售数据.xlsx",
            title="商城详细销售记录",
            data_desc=data_desc,
            sample_data=sample
        )


def fetch_finance_policies(self):
    """获取金融政策示例数据"""
    return [
        {
            "source": "https://finance.example.com/loan-policy",
            "content": """
                2024年第二季度信贷政策:

                1. 个人消费贷款:
                   - 利率: 年化7.2% (浮动)
                   - 期限: 6-36个月
                   - 额度: 最高50万元

                2. 房屋贷款:
                   - 首套房: LPR+20基点
                   - 二套房: LPR+60基点
                   - 贷款材料:
                     * 身份证复印件
                     * 6个月银行流水
                     * 收入证明

                3. 新用户优惠:
                   - 首贷利率优惠0.5%
                   - 有效期至2024-06-30
                """
        }
    ]


if __name__ == '__main__':
    doc = DocumentLoader()
    doc.fetch_ecommerce_docs()
